package com.tsh.chat.controller;

import com.fasterxml.jackson.core.JacksonException;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.tsh.chat.domain.ChatHistory;
import com.tsh.chat.domain.MessageBody;
import com.tsh.chat.domain.Personality;
import com.tsh.chat.service.ChatService;
import com.tsh.chat.service.PersonalitySimulatorService;
import com.tsh.chat.service.PromptService;
import com.tsh.chat.util.ChatUtil;
import com.tsh.common.core.utils.ServletUtils;
import com.tsh.common.core.web.controller.BaseController;
import com.tsh.common.core.web.domain.AjaxResult;
import com.tsh.common.security.utils.SecurityUtils;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageAccumulator;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.Choice;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/sse")
public class SseController extends BaseController {

    @Autowired
    private ChatService chatService;

    @Autowired
    private PersonalitySimulatorService personalitySimulatorService;

    private final Map<String, SseEmitter> sseEmitters = new ConcurrentHashMap<>();

    /**
     * 获取SseEmitter
     */
    @GetMapping(value = "/getStream/{userId}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter getConnection(@PathVariable String userId, HttpServletResponse response){
        //设置响应的内容类型为HTML，字符集为UTF-8。
        response.setContentType("text/html;charset=utf-8");
        // 设置响应头，指明这是一个EventSource响应
        response.setHeader("Content-Type", "text/event-stream");
        //设置响应头"Cache-Control"为"no-cache"，表示不允许缓存此响应。
        response.setHeader("Cache-Control", "no-cache");
        //设置响应头"Connection"为"keep-alive"，表示保持连接。
        response.setHeader("Connection", "keep-alive");

        SseEmitter sseEmitter = new SseEmitter(5*60*1000L);
        sseEmitters.put(userId, sseEmitter);
        sseEmitter.onCompletion(() -> {
            sseEmitters.remove(userId);
            logger.info("SseEmitter completed for user: " + userId);
        });
        sseEmitter.onTimeout(() -> {
            sseEmitters.remove(userId);
            logger.info("SseEmitter timed out for user: " + userId);
        });
        sseEmitter.onError((e) -> {
            sseEmitters.remove(userId);
            logger.error("SseEmitter error for user: " + userId, e);
        });

        logger.info("Created SseEmitter for user: " + userId);
        return sseEmitter;
    }

    @PostMapping("/stream")
    public AjaxResult chatStream(@RequestParam int userId, @RequestParam String message) {
        logger.info("Attempting to get SseEmitter for user: " + userId);
        String Uuid = String.valueOf(userId);
        SseEmitter emitter = sseEmitters.get(Uuid);
        if (emitter == null) {
            System.out.println("SseEmitter is null!");
        }

        List<ChatMessage> messages = new ArrayList<>();
        List<ChatHistory> chatHistoryList = chatService.getChatHistoryByUserId(userId);
        for (ChatHistory content: chatHistoryList){
            ChatMessage chatMessage = new ChatMessage(content.getRole(), content.getContent());
            messages.add(chatMessage);
        }
        ChatMessage newChatMessage = new ChatMessage(ChatMessageRole.USER.value(), message);
        messages.add(newChatMessage);
        try{
            String newAnswer = ChatUtil.chatStream(messages, emitter);
            /**
             * 获取最新回复成功，新增的对话保存到数据库
             */
            chatService.insertNewChat(userId, ChatMessageRole.USER.value(), message);
            chatService.insertNewChat(userId, ChatMessageRole.ASSISTANT.value(), newAnswer);
        }catch (JsonProcessingException e){
            e.printStackTrace();
            logger.info("连续对话模块出现了意外...");
            sseEmitters.remove(Uuid);
            return error();
        }
        if(emitter != null){
            emitter.complete();
        }
        sseEmitters.remove(Uuid);
        return success();
    }

    /**
     * 人设模拟对话流式输出
     */
    @PostMapping("/constantChat")
    public AjaxResult constantChat(@RequestBody MessageBody messageBody) {
        if(messageBody.getUuid() != SecurityUtils.getUserId()){
            return error("无权操作!请求ID:" + messageBody.getUuid() + "当前用户ID:" + SecurityUtils.getUserId());
        }
        if(messageBody.getMessage().isEmpty()) return error("内容不能为空!");
        String Uuid = String.valueOf(messageBody.getUuid());
        SseEmitter emitter = sseEmitters.get(Uuid); //获取emitter
        List<Personality> list = personalitySimulatorService.getPersonality(messageBody.getUuid());
        if(list.isEmpty()){
            if(emitter != null){
                emitter.complete();
            }
            sseEmitters.remove(Uuid);
            return warn("您还没有创建角色，请先添加人设信息");
        }else{
            Date now = new Date();
            try{
                String desc = list.get(0).getDescription();
                String achieve = list.get(0).getAchievement();
                String relation = list.get(0).getRelationship();
                List<ChatMessage> messages = new ArrayList<>();
                List<ChatHistory> chatHistoryList = personalitySimulatorService.getChatHistoryByUserId(messageBody.getUuid());
                for (ChatHistory content: chatHistoryList){
                    ChatMessage chatMessage = new ChatMessage(content.getRole(), content.getContent());
                    messages.add(chatMessage);
                }
                ChatMessage newChatMessage = new ChatMessage(ChatMessageRole.USER.value(), messageBody.getMessage());
                messages.add(newChatMessage);
                String newAnswer = ChatUtil.constantSimulatedChatSSE(messages, desc, achieve, relation, now, emitter);
                personalitySimulatorService.insertChatHistory(messageBody.getUuid(), ChatMessageRole.USER.value(), messageBody.getMessage());
                personalitySimulatorService.insertChatHistory(messageBody.getUuid(), ChatMessageRole.ASSISTANT.value(), newAnswer);
            }catch (JacksonException e){
                sseEmitters.remove(Uuid);
                return error(e.getMessage());
            }
            if(emitter != null){
                emitter.complete();
            }
            sseEmitters.remove(Uuid);
            return success();
        }
    }

}
